""" Full assembly of the parts to form the complete network """
import sys
import os
cur_path = os.path.split(__file__)[0]
sys.path.append(cur_path)
from unet_parts import *


class UNet(nn.Module):
    def __init__(self, in_dim, out_dim, bilinear=False):
        super(UNet, self).__init__()
        self.n_channels = in_dim
        self.out_dim = out_dim
        self.bilinear = bilinear
        self.reduce = 2
        self.inc = (DoubleConv(in_dim, 64 // self.reduce))
        self.down1 = (Down(64// self.reduce, 128// self.reduce))
        self.down2 = (Down(128// self.reduce, 256// self.reduce))
        self.down3 = (Down(256// self.reduce, 512// self.reduce))
        factor = 2 if bilinear else 1
        self.down4 = (Down(512// self.reduce, 1024 // factor// self.reduce))
        self.up1 = (Up(1024// self.reduce, 512 // factor// self.reduce, bilinear))
        self.up2 = (Up(512// self.reduce, 256 // factor// self.reduce, bilinear))
        self.up3 = (Up(256// self.reduce, 128 // factor// self.reduce, bilinear))
        self.up4 = (Up(128// self.reduce, 64// self.reduce, bilinear))
        self.outc = (OutConv(64 // self.reduce, out_dim))

    def forward(self, x):
        x1 = self.inc(x)
        x2 = self.down1(x1)
        x3 = self.down2(x2)
        x4 = self.down3(x3)
        x5 = self.down4(x4)
        x = self.up1(x5, x4)
        x = self.up2(x, x3)
        x = self.up3(x, x2)
        x = self.up4(x, x1)
        logits = self.outc(x)
        return logits

    def use_checkpointing(self):
        self.inc = torch.utils.checkpoint(self.inc)
        self.down1 = torch.utils.checkpoint(self.down1)
        self.down2 = torch.utils.checkpoint(self.down2)
        self.down3 = torch.utils.checkpoint(self.down3)
        self.down4 = torch.utils.checkpoint(self.down4)
        self.up1 = torch.utils.checkpoint(self.up1)
        self.up2 = torch.utils.checkpoint(self.up2)
        self.up3 = torch.utils.checkpoint(self.up3)
        self.up4 = torch.utils.checkpoint(self.up4)
        self.outc = torch.utils.checkpoint(self.outc)


# if __name__ == '__main__':


#     torch.cuda.set_device(0)
#     net =UNet(1, 35).cuda().eval()

#     data = torch.randn(4, 1, 64, 128, 64).cuda()

#     out = net(data)
#     print('?', out.shape)
